
__all__ = []

import os
import pyutilib.plugin
from coopr.opt.parallel.manager import *
from coopr.opt.parallel.solver import *
from coopr.opt.results import SolverResults
import kestrel

import pyutilib


class SolverManager_NEOS(AsynchronousSolverManager):

    def clear(self):
        """
        Clear manager state
        """
        AsynchronousSolverManager.clear(self)
        self.kestrel = kestrel.kestrelAMPL()
        self._opt = None
        self._ah = {} # index is the NEOS job number

        # to grab streamed output from NEOS, need to keep
        # map of action handle to the to-date string of
        # extracted output.
        # TBD: The following entries aren't currently cleaned up, but
        #      we're still trying to get the basics down.
        # store pairs of NEOS message offset and NEOS message string.
        # index into the map is the NEOS job number
        self._neos_log = {} 

    def _perform_queue(self, ah, *args, **kwds):
        """
        Perform the queue operation.  This method returns the ActionHandle,
        and the ActionHandle status indicates whether the queue was successful.
        """
        if 'opt' in kwds:
            self._opt = kwds['opt']
            del kwds['opt']
        else:
            raise ActionManagerError, "Undefined solver"
        self._opt._presolve(*args, **kwds)
        #
        # Apply kestrel
        #
        os.environ['kestrel_options'] = 'solver=%s' % self._opt.neos_name
        os.environ[self._opt.neos_name.lower()+'_options'] = self._opt.solver_options
        xml = self.kestrel.formXML(self._opt._problem_files[0])
        (jobNumber, password) = self.kestrel.submit(xml)
        ah.job = jobNumber
        ah.password = password
        #
        # Store action handle, and return
        #
        self._ah[jobNumber] = ah
        self._neos_log[jobNumber] = (0, "")
        return ah

    def _perform_wait_any(self):
        """
        Perform the wait_any operation.  This method returns an
        ActionHandle with the results of waiting.  If None is returned
        then the ActionManager assumes that it can call this method again.
        Note that an ActionHandle can be returned with a dummy value,
        to indicate an error.
        """
        for jobNumber in self._ah:

            status = self.kestrel.neos.getJobStatus(jobNumber,self._ah[jobNumber].password)

            if not status in ("Running", "Waiting"):

                ah = self._ah[jobNumber]                
                
                # the job is done.
                self._ah[jobNumber] = None
                ah.status = ActionStatus.done
                
                # retrieve the final results, which are in message/log format.
                results = self.kestrel.neos.getFinalResults(jobNumber, ah.password)

                (current_offset, current_message) = self._neos_log[jobNumber]
                OUTPUT=open(self._opt.log_file,'w')
                print >>OUTPUT, current_message
                OUTPUT.close()

                OUTPUT=open(self._opt.soln_file,'w')
                print >>OUTPUT, results.data
                OUTPUT.close()

                rc = None
                solver_results = self._opt.process_output(rc)
                solver_results.symbol_map = self._opt._symbol_map
                self.results[ah.id] = solver_results
                return ah
            else:
                # grab the partial messages from NEOS as you go, in case you want
                # to output on-the-fly. we don't currently do this, but the infrastructure
                # is in place.
                (current_offset, current_message) = self._neos_log[jobNumber]
                # TBD: blocking isn't the way to go, but non-blocking was triggering some exception in kestrel.
                (message_fragment, new_offset) = self.kestrel.neos.getIntermediateResults(jobNumber, self._ah[jobNumber].password, current_offset)
                print message_fragment,
                self._neos_log[jobNumber] = (new_offset, current_message + message_fragment.data)

        return None


SolverManagerRegistration("neos", SolverManager_NEOS)

